#include "transcoder_core.h"
extern int running_enc_t_cnt;
enum AVPixelFormat src_pix_fmt;
int src_w, src_h;
int gop_size;
int src_image_size;
int running_enc_t_cnt;
sem_t encoding_mutex;
int running_enc_thread_cnt;
int ENC_THREAD_CNT;
int PRINT_DYNAMIC_FPS;

AVRational CODEC_TIME_BASE;
AVRational STREAM_TIME_BASE;
int TICKS_PER_FRAME;

int GOP_COUNT;
int gop_size;
EncodeStruct ess[6];
int DEC_THREAD_CNT;
int ENC_THREAD_CNT;
int OUTPUT_VIDEO_NUM;

List **FULL_GOP_BUFFER;

pthread_mutex_t running_enc_thread_cnt_lock = PTHREAD_MUTEX_INITIALIZER;
extern x264_control_t *x264_ctrl;
int inject;

static inline int open_encoder (codec_t *c);
static inline int close_encoder(codec_t *c);

EncodeStruct *slowest_es()
{
    int min_cnt = ess[0].encoded_frame_cnt;
    EncodeStruct *slowest_es = &ess[0];

    for (int i = 1; i < OUTPUT_VIDEO_NUM; i++)
    {
        if (!ess[i].finished && ess[i].encoded_frame_cnt < min_cnt)
        {
            slowest_es = &ess[i];
            min_cnt = slowest_es->encoded_frame_cnt;
        }
    } 

    return slowest_es;
}

int all_finished()
{
    for (int i = 0; i < OUTPUT_VIDEO_NUM; i++)
        if (!ess[i].finished)
            return 0;

    return 1;
}

static inline int open_encoder (codec_t *c)
{
    /*set related X264 parameters*/
    AVDictionary *param = NULL;
    if(c->codec_ctx->codec_id == AV_CODEC_ID_H264)
    {
        av_dict_set(&param, "tune", "zerolatency", AV_DICT_APPEND);
        /*av_dict_set(&param,"refs","3", AV_DICT_APPEND);*/
        /*av_dict_set(&param,"preset","medium",0);*/
        /*av_dict_set(&param,"preset","veryslow",0);*/
    }
    else if(c->codec_ctx->codec_id == AV_CODEC_ID_H265)
    {
        av_dict_set(&param, "tune", "zerolatency", AV_DICT_APPEND);
        av_dict_set(&param, "x265-params", "pools=none", AV_DICT_APPEND);
        av_dict_set(&param, "x265-params", "frame-threads=1", AV_DICT_APPEND);
    }
    else
    {
        printf("Encode format Error!\n");
        exit(1);
    }

    /*open encoder*/
    if (avcodec_open2(c->codec_ctx, c->codec, &param) < 0)
    /*if (avcodec_open2(c->codec_ctx, c->codec, NULL) < 0)*/
    {
		printf("%s Failed to open encoder! \n", "scale and encode ");
		return -1;
	}

    /*x264 library*/
    if (inject)
        c->x264_ctrl->call_cnt += 1;

    return 0;
}

static inline int close_encoder(codec_t *c)
{
	avcodec_close(c->video_st->codec);

    /*x264 library*/
    if (inject)
        c->x264_ctrl->call_cnt++;

    return 0;
}

int init_codec_ctx (codec_t *c, EncodeStruct *es)
{
    /*1. allocate AVFormatContext*/
    c->fmt_ctx = avformat_alloc_context();

    /*2. allocate AVOutputFormat*/
    char guess_file[100];
	sprintf(guess_file, "0.%s", es->output_format);
    c->out_fmt = av_guess_format(NULL, guess_file, NULL);
    c->fmt_ctx->oformat = c->out_fmt;

    /*3. allocate AVStream*/
	c->video_st = avformat_new_stream(c->fmt_ctx, 0);
	if (c->video_st==NULL){
		printf("Failed to new video stream! \n");
		return -1;
	}

    /*4. initialize AVCodecContext*/
	c->codec_ctx = c->video_st->codec;
	c->video_st->time_base = STREAM_TIME_BASE;

	c->codec_ctx->codec_id = c->out_fmt->video_codec;
	c->codec_ctx->codec_type = AVMEDIA_TYPE_VIDEO;
	c->codec_ctx->pix_fmt = es->dst_pix_fmt;
	c->codec_ctx->width = es->dst_width;
	c->codec_ctx->height = es->dst_height;

	c->codec_ctx->time_base = CODEC_TIME_BASE;
	c->codec_ctx->ticks_per_frame = TICKS_PER_FRAME;

	c->codec_ctx->bit_rate = es->bitrate;
	c->codec_ctx->rc_max_rate = es->bitrate;
	c->codec_ctx->rc_min_rate = es->bitrate;
	c->codec_ctx->rc_buffer_size = es->bitrate;
	/*c->codec_ctx->bit_rate_tolerance = 0;*/
    /*c->codec_ctx->bit_rate_tolerance = 4000000;*/

    /*c->codec_ctx->gop_size = 5;*/
    c->codec_ctx->gop_size = 3;
    /*c->codec_ctx->gop_size = 4;*/
    /*c->codec_ctx->gop_size = -1;*/
    if (es->dst_width > 640)
    {
        c->codec_ctx->gop_size = 4;
    }
    else
    {
        c->codec_ctx->gop_size = 10;
    }

    /*c->codec_ctx->qmin = 10;*/
    c->codec_ctx->qmin = -1;
    c->codec_ctx->qmax = -1;
    /*c->codec_ctx->qmax = 51;*/
    /*c->codec_ctx->max_b_frames = 0;*/
    /*c->codec_ctx->max_b_frames = 3;*/
    /*c->codec_ctx->max_qdiff = 4;*/
    c->codec_ctx->thread_count = 1;
    /*printf("%d \n", c->codec_ctx->gop_size);*/
	// c->codec_ctx->thread_count = 60;
	/*c->codec_ctx->trellis = 1;*/
    c->codec_ctx->trellis = -1;
    /*c->codec_ctx->me_range = 16;*/
	/*c->codec_ctx->me_method=ME_HEX;*/
    /*c->codec_ctx->me_method=-1;*/
	c->codec_ctx->coder_type = 1;
	/*c->codec_ctx->me_cmp += FF_CMP_CHROMA;*/
    /*c->codec_ctx->me_cmp = -1;*/
	/*c->codec_ctx->me_subpel_quality = 7;*/
    /*c->codec_ctx->me_subpel_quality = -1;*/
    /*c->codec_ctx->max_b_frames = 0;*/

	// c->codec_ctx->flags &= ~CODEC_FLAG_QSCALE;
	// c->profile = FF_PROFILE_H264_BASELINE;

    /*find encoder*/
	c->codec = avcodec_find_encoder(c->codec_ctx->codec_id);
	if (!c->codec)
    {
		printf(" %s Can not find encoder! \n", "scale and encode ");
		return -1;
	}

    /*create sws_scale context*/
	c->sws_ctx = sws_getContext(src_w, src_h, src_pix_fmt,
                                            es->dst_width,
                                            es->dst_height,
                                            es->dst_pix_fmt,
			                                SWS_BILINEAR, NULL, NULL, NULL);

	if (!c->sws_ctx) {
		printf("Error 33: fail to create scale context for the conversion "
				"fmt:%s s:%dx%d -> fmt:%s s:%dx%d\n",
				av_get_pix_fmt_name(src_pix_fmt), src_w, src_h,
				av_get_pix_fmt_name(es->dst_pix_fmt), es->dst_width, es->dst_height);

		exit(1);
	}

    /*7.
     *to transcode, we need src frame picture and dst frame and scale context,
     *in this version, every GOP doesn't store the whole AVframe, but the
     *decompressed data, so we need to reconstruct the src frame
     */
	c->src_frame = av_frame_alloc();
	c->src_frame->width = src_w;
	c->src_frame->height = src_h;
	c->src_frame->format = src_pix_fmt;

	c->dst_frame = av_frame_alloc();
	c->dst_frame->width = es->dst_width;
	c->dst_frame->height = es->dst_height;
	c->dst_frame->format = es->dst_pix_fmt;

    /*fill the dst_frame with image buffer*/
    int image_size = av_image_get_buffer_size(es->dst_pix_fmt, es->dst_width,
                                                es->dst_height, 1);

    uint8_t *image_buf = (uint8_t *) av_malloc(image_size);
    av_image_fill_arrays(c->dst_frame->data,
                        c->dst_frame->linesize,
                        image_buf,
                        es->dst_pix_fmt,
                        es->dst_width,
                        es->dst_height,
                        1);

    /*8. allocate a new AVPacket*/
    av_init_packet(&c->pkt);

    return 0;
}

int close_codec_ctx (codec_t *c)
{
    av_free(c->dst_frame);
    av_free(c->src_frame);
	sws_freeContext(c->sws_ctx);
	avformat_free_context(c->fmt_ctx);

    return 0;
}

int scale_and_encode_gop(Gop *g, EncodeStruct *es, codec_t *c)
{
    /*x264 library*/
    if (inject)
        c->x264_ctrl->call_cnt = 0;

    open_encoder(c);

    /*generate these gop output filename*/
	char gop_output[100];
	sprintf(gop_output, "./%s/%d.%s",
            es->output_dir, g->gop_id, es->output_format);

    if (avio_open(&c->fmt_ctx->pb, gop_output, AVIO_FLAG_READ_WRITE) < 0)
    {
        printf("fail to open output file\n");
        return -1;
    }

    if (avformat_write_header(c->fmt_ctx, NULL) < 0)
    {
        printf("fail to open output file header\n");
        return -1;
    }

    AVFrame *src_frame = c->src_frame;
    AVFrame *dst_frame = c->dst_frame;

    /*Now every initialization has been set up*/
	for (int i = 0; i < gop_size; i++)
    {
		if(g->valid[i] == 0)
            continue;

        /*fetch the pixel buffer and fill it in src_frame*/
        av_image_fill_arrays(src_frame->data,
                            src_frame->linesize,
                            g->buf + i * src_image_size,
                            src_pix_fmt,
                            src_w,
                            src_h,
                            1);

        /*if (i % 1 == 0) {*/
            /*dst_frame->pict_type = AV_PICTURE_TYPE_I;*/
            /*src_frame->pict_type = AV_PICTURE_TYPE_I;*/
        /*} else {*/
            /*dst_frame->pict_type = AV_PICTURE_TYPE_NONE;*/
            /*src_frame->pict_type = AV_PICTURE_TYPE_NONE;*/
        /*}*/

        /*
         *only if the resolution of src_frame and dst_frame is different do
         *we need scale
         */

        if(es->dst_width != src_w || es->dst_height != src_h){
            sws_scale(c->sws_ctx,
                    (const uint8_t **)src_frame->data, src_frame->linesize,
                    0, src_h,
                    dst_frame->data, dst_frame->linesize);

        } else {
            dst_frame = src_frame;
        }

        dst_frame->pts = g->gop_id * gop_size + i;

        /*invoke the ffmpeg's encoding function*/
        int got_picture;
        c->pkt.data = NULL;
        c->pkt.size = 0;

        /*x264 library*/
        if (inject)
            c->x264_ctrl->call_cnt = 0;

        if (avcodec_encode_video2(c->codec_ctx, &c->pkt, dst_frame, &got_picture) < 0)
        {
            printf("fail to encode frame\n");
        }

		if (got_picture)
        {
            /*printf("%d\n", c->pkt.flags);*/
			c->pkt.stream_index = c->video_st->index;
            av_write_frame(c->fmt_ctx, &c->pkt);
		}

        atomic_inc(&es->encoded_frame_cnt, 1, &es->encoded_frame_mutex);
	}

    /*flush encoder*/
    int ret = flush_encoder(c, es, 0);
    if (ret < 0) {
        printf("Flushing encoder failed\n");
        return -1;
    }

    av_write_trailer(c->fmt_ctx);
    avio_close(c->fmt_ctx->pb);

    /*x264 library*/
    if (inject)
        c->x264_ctrl->call_cnt = 0;

    close_encoder(c);

    return 0;
}

int encode_and_scale_thread(void *arg)
{
	binding_core();
    int flag = 0;

    /*x264_library*/
    /*
     *inject = 1;
     *set_injected_version(1);
     */

    /*get the encode struct this thread will serve*/
	EncodeStruct * es = (EncodeStruct *)arg;

reencode:

    atomic_inc(&es->running_thread_cnt, 1, &es->thread_cnt_mutex);
    int full_gop_buffer_id = es->full_gop_buffer_id;

    codec_t codec_var;

    /*get the x264_control_t*/
    x264_control_t *ct;
    if (inject)
    {
        ct = get_x264_ctrl(1);
        if(ct == NULL)
        {
            printf("transcoder fail to get an x264_control_t\n");
            return -1;
        }

        codec_var.x264_ctrl = ct;
    }

    /*initialize encoding parameters, like AVFormatContext ...*/
    init_codec_ctx(&codec_var, es);

    /*using a loop to get the gop from the queue and encode the gop*/
    while (true)
    {
	    Gop *g = remove_gop_from_full_gop_buffer(full_gop_buffer_id);

        if (flag == 0)
        {
            /*only after all threads have been launched, will the average fps be printed*/
            pthread_mutex_lock(&running_enc_thread_cnt_lock);
	        if(running_enc_thread_cnt < ENC_THREAD_CNT)
                ++running_enc_thread_cnt;
            pthread_mutex_unlock(&running_enc_thread_cnt_lock);

            flag = 1;

            if (running_enc_thread_cnt == ENC_THREAD_CNT)
                PRINT_DYNAMIC_FPS = 1;

        }
        
        if (is_empty_Gop(g))
        {
	        add_gop_to_full_gop_buffer(g, full_gop_buffer_id);
            break;
        }

        scale_and_encode_gop(g, es, &codec_var);
        free_Gop(g);
    }

    close_codec_ctx(&codec_var);

    es->finished = 1;

    /*free the x264_control_t*/
    if (inject)
        free_ctrl(ct);

    atomic_dec(&es->running_thread_cnt, 1, &es->thread_cnt_mutex);

    /*output the video if all the threads in these es exit*/
    if (es->running_thread_cnt == 0)
    {
        output_video(es);
    } 

    /*one task has finished, it should check whether to encode other output*/
    if (!all_finished())
    {
        es = slowest_es();
        goto reencode;
    }

	PRINT_DYNAMIC_FPS = 0;

	return 0;
}

int multi_encode_thread(void *arg)
{
	/*
     * georgefan 13:40/9.9
     * set affinity CPU
     */
    binding_core();

    /*x264_library*/
    inject = 1;
    set_injected_version(1);

    /*only after all threads have been launched, will the average fps be printed*/
    pthread_mutex_lock(&running_enc_thread_cnt_lock);
	if(running_enc_thread_cnt < ENC_THREAD_CNT)
        ++running_enc_thread_cnt;
    pthread_mutex_unlock(&running_enc_thread_cnt_lock);

    /*get the encode struct this thread will serve*/
	EncodeStructList *es_list = (EncodeStructList *)arg;
    int es_num = es_list->es_num;

    /*get the x264_control_t*/
    x264_control_t *ct = get_x264_ctrl(es_num);
    if(ct == NULL)
    {
        printf("transcoder fail to get an x264_control_t\n");
        return -1;
    }

    /*
     * fsh13:58/9.9
     * declare and initization of codec_t
     */
    codec_t *codec_t_list = (codec_t *)malloc(es_num * sizeof(codec_t)); 

    for (int i = 0; i < es_num; i++)
    {
        codec_t_list[i].x264_ctrl = ct;
        atomic_inc(&es_list->es[i]->running_thread_cnt, 1,
                &es_list->es[i]->thread_cnt_mutex);

        init_codec_ctx(&codec_t_list[i], es_list->es[i]);
    }

    if (running_enc_thread_cnt == ENC_THREAD_CNT)
        PRINT_DYNAMIC_FPS = 1;

    /*using a loop to get the gop from the queue and encode the gop*/
    while (true)
    {
        Gop *g ;

        for (int i = es_num - 1; i >= 0; i--)
        {
	        g = remove_gop_from_full_gop_buffer(es_list->es[i]->full_gop_buffer_id);
                
            if (is_empty_Gop(g))
            {
	            add_gop_to_full_gop_buffer(g, es_list->es[i]->full_gop_buffer_id);
                goto finish;
            }
        }
        
        gop_multi_scale_enc(g, es_list, codec_t_list);

        for (int i = 0; i < es_num; i++)
            free_Gop(g);
    }

finish:

	PRINT_DYNAMIC_FPS = 0;

    for (int i = 0; i < es_num; i++)
    {
        atomic_dec(&es_list->es[i]->running_thread_cnt, 1,
                &es_list->es[i]->thread_cnt_mutex);

        close_codec_ctx(&codec_t_list[i]);

        /*output the video if all the threads in these es exit*/
        if (es_list->es[i]->running_thread_cnt == 0)
            output_video(es_list->es[i]);
    }

	return 0;
}

int flush_multi_encoder (codec_t *c_list, unsigned int stream_idx, int num)
{
    codec_t *c;

    for (int i = 0; i < num; i++)
    {
	    if (!(c_list[0].fmt_ctx->streams[stream_idx]->codec->codec->capabilities &
                    CODEC_CAP_DELAY))
            return 0;
    }

    int ret;
    int got_frame;

    while (1)
    {
        for (int i = 0; i < num; i++)
        {
            c = &c_list[i];
            c->pkt.data = NULL;
            c->pkt.size = 0;
            c->x264_ctrl->call_cnt = i; 
            c->x264_ctrl->skip_call = i;

		    ret = avcodec_encode_video2 (c->fmt_ctx->streams[stream_idx]->codec,
                    &c->pkt,
		    		NULL,
                    &got_frame);

            if (ret < 0)
                break;

            if (!got_frame)
            {
                ret = 0;
                break;
            }

            ret = av_write_frame(c->fmt_ctx, &c->pkt);
        }
    }

    return ret;
}

int flush_encoder (codec_t *c, EncodeStruct *es, unsigned int stream_idx)
{
	if (!(c->fmt_ctx->streams[stream_idx]->codec->codec->capabilities & CODEC_CAP_DELAY))
		return 0;
    
    int ret;
    int got_frame;

    while (1)
    {
        c->pkt.data = NULL;
        c->pkt.size = 0;

        /*x264 library*/
        if (inject)
            c->x264_ctrl->call_cnt = 0; 

		ret = avcodec_encode_video2 (c->fmt_ctx->streams[stream_idx]->codec,
                &c->pkt, NULL, &got_frame);

        if (ret < 0)
            break;

        if (!got_frame)
        {
            ret = 0;
            break;
        }

        /*printf("%d\n", c->pkt.flags);*/
        ret = av_write_frame(c->fmt_ctx, &c->pkt);
    }

    return ret;
}

int gop_multi_scale_enc(Gop *g, EncodeStructList *es_list, codec_t *c_list)
{
    int es_num = es_list->es_num;
    
    codec_t *c;
    EncodeStruct *es;

    /*x264 library*/
    if (inject)
        c_list[0].x264_ctrl->call_cnt = 0;

    /*we must initialize codec for each output*/
    for (int i = 0; i < es_num; i++)
    {
        c = &c_list[i];
        es = es_list->es[i];

        open_encoder(c);

        /*generate these gop output filename*/
        char gop_output[100];
        sprintf(gop_output, "./%s/%d.%s",
                es->output_dir, g->gop_id, es->output_format); 
        if (avio_open(&c->fmt_ctx->pb, gop_output, AVIO_FLAG_READ_WRITE) < 0)
        {
            printf("fail to open output file\n");
            return -1;
        }

        if (avformat_write_header(c->fmt_ctx, NULL) < 0)
        {
            printf("fail to open output file header\n");
            return -1;
        }
    }

    AVFrame *src_frame = c_list[0].src_frame;
    AVFrame *dst_frame = c_list[0].dst_frame;
    es = es_list->es[0];

        /*dst_frame->pict_type = AV_PICTURE_TYPE_I;*/
        /*src_frame->pict_type = AV_PICTURE_TYPE_I;*/
    /*Now every initialization has been set up*/
	for (int i = 0; i < gop_size; i++)
    {
		if(g->valid[i] == 0)
            continue;

        /*fetch the pixel buffer and fill it in src_frame*/
        av_image_fill_arrays(src_frame->data,
                            src_frame->linesize,
                            g->buf + i * src_image_size,
                            src_pix_fmt,
                            src_w,
                            src_h,
                            1);


        /*invoke the ffmpeg's encoding function*/
        for (int j = 0; j < es_num; j++)
        {
            c = &c_list[j];
            c->pkt.data = NULL;
            c->pkt.size = 0;
            c->x264_ctrl->call_cnt = j; 
            /*c->x264_ctrl->call_cnt = 0; */
            c->x264_ctrl->skip_call = j;

     /*       es = es_list->es[j];*/
            /*AVFrame *dst_frame = c->dst_frame;*/
            /*
             *only if the resolution of src_frame and dst_frame is different do
             *we need scale
             */
            if(es->dst_width != src_w || es->dst_height != src_h){
                sws_scale(c->sws_ctx,
                        (const uint8_t **)src_frame->data, src_frame->linesize,
                        0, src_h,
                        dst_frame->data, dst_frame->linesize);

            } else {
                dst_frame = src_frame;
            }

            dst_frame->pts = g->gop_id * gop_size + i;

            int got_picture;

            if (avcodec_encode_video2(c->codec_ctx, &c->pkt, dst_frame, &got_picture) < 0)
            {
                printf("faile to encode frame\n");
            }

		    if (got_picture)
            {
		    	c->pkt.stream_index = c->video_st->index;
                /*printf("%d type\n", c->codec_ctx->coded_frame->pict_type);*/
                av_write_frame(c->fmt_ctx, &c->pkt);
		    }

            atomic_inc(&es->encoded_frame_cnt, 1, &es->encoded_frame_mutex);
        }
	}


    /*flush encoder*/
    /*
	 *int ret = flush_multi_encoder(c_list, 0, es_num);
	 *if (ret < 0) {
	 *    printf("Flushing encoder failed\n");
	 *    return -1;
	 *}
     */

    /*x264 library*/
    if (inject)
        c_list[0].x264_ctrl->call_cnt = 0;

    for (int i = 0; i < es_num; i++)
    {
        c = &c_list[i];
	    av_write_trailer(c->fmt_ctx);
	    avio_close(c->fmt_ctx->pb);

        close_encoder(c);
    }
    
    return 0;
}

int initialize_encode_data()
{
	sem_init(&encoding_mutex, 0, 1);

	int i;
	for(i = 0; i < OUTPUT_VIDEO_NUM; i++)
    {
		ess[i].running_thread_cnt = 0;
		sem_init(&ess[i].thread_cnt_mutex, 0, 1);
		sem_init(&ess[i].encoded_frame_mutex, 0, 1);
	}

	return 0;
}
